#include "kalman_filter.h"

// 卡尔曼滤波器创建
KalmanFilter *create_kalman_filter(Matrix *A, Matrix *H)
{
    KalmanFilter *kf = (KalmanFilter *)malloc(sizeof(KalmanFilter));
    kf->A = matrix_copy(A);
    kf->H = matrix_copy(H);

    // 初始化噪声矩阵
    kf->Q = matrix_eye(A->rows);
    kf->R = matrix_eye(H->rows);
    kf->P = matrix_eye(A->rows);
    kf->X = create_matrix(A->rows, 1);
    kf->X_pre = create_matrix(A->rows, 1);
    kf->P_pre = matrix_eye(A->rows);

    return kf;
}

// 卡尔曼滤波器释放
void free_kalman_filter(KalmanFilter *kf)
{
    if (kf)
    {
        free_matrix(kf->A);
        free_matrix(kf->H);
        free_matrix(kf->Q);
        free_matrix(kf->R);
        free_matrix(kf->P);
        free_matrix(kf->X);
        free_matrix(kf->X_pre);
        free_matrix(kf->P_pre);
        free(kf);
    }
}

// 卡尔曼预测
void kalman_predict(KalmanFilter *kf)
{
    // X_pre = A * X
    Matrix *temp = matrix_multiply(kf->A, kf->X);
    memcpy(kf->X_pre->data, temp->data, kf->X_pre->rows * kf->X_pre->cols * sizeof(double));
    free_matrix(temp);

    // P_pre = A * P * A^T + Q
    temp = matrix_multiply(kf->A, kf->P);
    Matrix *temp2 = matrix_transpose(kf->A);
    Matrix *temp3 = matrix_multiply(temp, temp2);
    free_matrix(temp);
    free_matrix(temp2);

    temp = matrix_add(temp3, kf->Q);
    memcpy(kf->P_pre->data, temp->data, kf->P_pre->rows * kf->P_pre->cols * sizeof(double));
    free_matrix(temp3);
    free_matrix(temp);
}

// 卡尔曼更新
void kalman_update(KalmanFilter *kf, Matrix *Z)
{
    // 计算残差 D = Z - H * X_pre
    Matrix *temp = matrix_multiply(kf->H, kf->X_pre);
    Matrix *D = matrix_subtract(Z, temp);
    free_matrix(temp);

    // 计算S = H * P_pre * H^T + R
    temp = matrix_multiply(kf->H, kf->P_pre);
    Matrix *temp2 = matrix_transpose(kf->H);
    Matrix *temp3 = matrix_multiply(temp, temp2);
    free_matrix(temp);
    free_matrix(temp2);

    Matrix *S = matrix_add(temp3, kf->R);
    free_matrix(temp3);

    // 计算卡尔曼增益 K = P_pre * H^T * S^{-1}
    temp2 = matrix_transpose(kf->H);
    temp = matrix_multiply(kf->P_pre, temp2);
    free_matrix(temp2);

    Matrix *S_inv = create_matrix(S->rows, S->cols);
    matrix_inverse(S, S_inv);

    Matrix *K = matrix_multiply(temp, S_inv);
    free_matrix(temp);
    free_matrix(S_inv);

    // 更新状态 X = X_pre + K * D
    temp = matrix_multiply(K, D);
    Matrix *X_new = matrix_add(kf->X_pre, temp);
    memcpy(kf->X->data, X_new->data, kf->X->rows * kf->X->cols * sizeof(double));
    free_matrix(temp);
    free_matrix(X_new);

    // 更新协方差 P = (I - K * H) * P_pre
    temp = matrix_multiply(K, kf->H);
    Matrix *I = matrix_eye(temp->rows);
    Matrix *temp4 = matrix_subtract(I, temp);
    free_matrix(I);
    free_matrix(temp);

    Matrix *P_new = matrix_multiply(temp4, kf->P_pre);
    memcpy(kf->P->data, P_new->data, kf->P->rows * kf->P->cols * sizeof(double));
    free_matrix(temp4);
    free_matrix(P_new);
    free_matrix(D);
    free_matrix(S);
    free_matrix(K);
}

// 卡尔曼滤波
void kalman_filt(KalmanFilter *kf, Matrix *Z)
{
    kalman_predict(kf);
    kalman_update(kf, Z);
}